package algorithm_demo.demo04.graph;

import java.util.Collection;
import java.util.Objects;
import java.util.Map;
import java.util.HashMap;
import java.util.Stack;

/**
 * 并查集代码
 *
 * @author Api
 * @date 2023/2/19 17:13
 */
public class UnionFind<V> {
    public static class Node<V> {
        public V value;

        public Node(V value) {
            this.value = value;
        }
    }

    //并查集三个Map
    public Map<V, Node<V>> nodes;

    public Map<Node<V>, Node<V>> parents;

    public Map<Node<V>, Integer> sizeMap;

    public UnionFind(Collection<V> values) {
        nodes = new HashMap<>();
        parents = new HashMap<>();
        sizeMap = new HashMap<>();
        for (V value : values) {
            Node<V> node = new Node<>(value);
            nodes.put(value, node);
            //自己指向自己
            parents.put(node, node);
            //记录数量,初始都是1
            sizeMap.put(node, 1);
        }
    }

    public Node<V> findFather(Node<V> cur) {
        Stack<Node<V>> path = new Stack<>();
        //需要校验当前节点是不是父节点
        while (cur != parents.get(cur)) {
            path.push(cur);
            cur = parents.get(cur);
        }
        //找到头节点，并且进行扁平化处理
        while (!path.isEmpty()) {
            //将各个节点均从parents中设置为头节点
            parents.put(path.pop(), cur);
        }
        return cur;
    }

    public boolean isSameSet(V a, V b) {
        if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
            return false;
        }
        return findFather(nodes.get(a)) == findFather(nodes.get(b));
    }

    public void unionSet(V a, V b) {
        //校验是否被注册
        if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
            return;
        }
        //查找a和b的头节点
        Node<V> aHead = findFather(nodes.get(a));
        Node<V> bHead = findFather(nodes.get(b));
        //如果aHead和bHead的内存地址相等直接pass
        if (aHead != bHead) {
            //查找aHead和bHead的长度，小挂大，这里是aHead大，bHead小
            int aSetSize = sizeMap.get(aHead);
            int bSetSize = sizeMap.get(bHead);
            if (aSetSize >= bSetSize) {
                parents.put(bHead, aHead);
                sizeMap.put(aHead, aSetSize+bSetSize);
                sizeMap.remove(bHead);
            }else{
                //aHead小，bHead大
                parents.put(aHead, bHead);
                sizeMap.put(bHead, aSetSize+bSetSize);
                sizeMap.remove(aHead);
            }
        }
    }
    public int getSetNums(){
        return sizeMap.size();
    }
}
